#include "CVehHook2.h"
#include <TlHelp32.h>
#include <Psapi.h>

std::map<ULONG_PTR, ULONG_PTR> CVehHook2::m_VehHook2Map;

CVehHook2::CVehHook2()
{
    m_lpVeh = AddVectoredExceptionHandler(1, &CVehHook2::VectoredExceptionHandler);
    if (!m_lpVeh)
        throw "Could not add Vectored Exception Handler";
}

CVehHook2::~CVehHook2()
{
    if (m_lpVeh)
        RemoveVectoredExceptionHandler(m_lpVeh);
}

LONG NTAPI CVehHook2::VectoredExceptionHandler(IN PEXCEPTION_POINTERS lpExceptionPointers)
{
#ifdef _WIN64
#define XIP Rip
#else
#define XIP Eip
#endif

    auto it = m_VehHook2Map.find((ULONG_PTR)lpExceptionPointers->ExceptionRecord->ExceptionAddress);
    if (it != m_VehHook2Map.end())
    {
        lpExceptionPointers->ContextRecord->XIP = it->second;
        return EXCEPTION_CONTINUE_EXECUTION;
    }

    lpExceptionPointers->ContextRecord->Dr0 = (ULONG_PTR)lpExceptionPointers->ExceptionRecord->ExceptionAddress;
    lpExceptionPointers->ContextRecord->Dr7 = 0x405;

    return EXCEPTION_CONTINUE_SEARCH;
}

BOOL CVehHook2::AddVehHook2(IN ULONG_PTR SrcFunc, IN ULONG_PTR NewFunc, IN DWORD dwId)
{
    HANDLE hSnapShot = INVALID_HANDLE_VALUE;
    THREADENTRY32 te32 = { 0 };
    BOOL bRet = FALSE;
    HANDLE hThread;
    CONTEXT ThreadCtx;

    do
    {
        if (!SrcFunc || !NewFunc || !dwId) break;

        auto it = m_VehHook2Map.find(SrcFunc);
        if (it != m_VehHook2Map.end())
        {
            hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, dwId);
            if (hThread)
            {
                ThreadCtx.ContextFlags = CONTEXT_ALL;
                bRet = GetThreadContext(hThread, &ThreadCtx);
                if (bRet)
                {
                    ThreadCtx.Dr0 = SrcFunc;
                    ThreadCtx.Dr7 = 0x1;
                    SetThreadContext(hThread, &ThreadCtx);
                    bRet = TRUE;
                }
                CloseHandle(hThread);
            }
            return bRet;
        }

        hSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, dwId);
        if (INVALID_HANDLE_VALUE == hSnapShot) break;

        te32.dwSize = sizeof(THREADENTRY32);
        if (!Thread32First(hSnapShot, &te32)) break;

        do
        {
            if (te32.th32OwnerProcessID == dwId)
            {
                hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, te32.th32ThreadID);
                if (hThread)
                {
                    ThreadCtx.ContextFlags = CONTEXT_ALL;
                    bRet = GetThreadContext(hThread, &ThreadCtx);
                    if (bRet)
                    {
                        ThreadCtx.Dr0 = SrcFunc;
                        ThreadCtx.Dr7 = 0x1;
                        SetThreadContext(hThread, &ThreadCtx);
                        m_VehHook2Map[SrcFunc] = NewFunc;
                        bRet = TRUE;
                    }

                    CloseHandle(hThread);
                }
            }

        } while (Thread32Next(hSnapShot, &te32));

    } while (FALSE);

    if (INVALID_HANDLE_VALUE != hSnapShot)
        CloseHandle(hSnapShot);

    return bRet;
}

BOOL CVehHook2::Remove(IN ULONG_PTR SrcFunc)
{
    CONTEXT ThreadCtx;
    HANDLE hThread;
    BOOL bRet;

    hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, GetCurrentThreadId());
    if (hThread)
    {
        ThreadCtx.ContextFlags = CONTEXT_ALL;
        bRet = GetThreadContext(hThread, &ThreadCtx);
        if (bRet)
        {
            ThreadCtx.Dr0 = 0;
            ThreadCtx.Dr7 = 0;
            SetThreadContext(hThread, &ThreadCtx);
            bRet = TRUE;
        }

        CloseHandle(hThread);
    }
    m_VehHook2Map.clear();
    return 0;
}
